import pandas as pd
import pymysql

from base import Config, logger

class MySQLClient:
	def __init__(self):
		# 初始化日志
		self.logger = logger
		try:
			# 连接 MySQL 数据库
			self.connection = pymysql.connect(
				host=Config().MYSQL_HOST,
				password=Config().MYSQL_PASSWORD,
				user=Config().MYSQL_USER,
				database=Config().MYSQL_DATABASE
			)
			# 创建游标
			self.cursor = self.connection.cursor()
			# 记录连接成功
			self.logger.info("MySQL 连接成功")
		except pymysql.MySQLError as e:
			# 记录连接失败
			self.logger.error(f"MySQL 连接失败：{e}")
			raise

	def create_table(self):
		create_table_query = """
								CREATE TABLE IF NOT EXISTS JPKB
								(
									id INT AUTO_INCREMENT PRIMARY KEY,
									subject_name VARCHAR(20),
									question VARCHAR(1000),
									answer VARCHAR(1000)
								)
							 """
		try:
			self.cursor.execute(create_table_query)
			self.connection.commit()
			self.logger.info("表创建成功")
		except pymysql.MySQLError as e:
			self.logger.error(f"表创建失败：{e}")
			raise

	def insert_data(self, csv_path):
		try:
			# 读取csv文件
			df = pd.read_csv(csv_path)
			for _, row in df.iterrows():
				insert_query = "INSERT INTO JPKB (subject_name, question, answer) VALUES (%s, %s, %s)"
				self.cursor.execute(insert_query, (row['学科名称'], row['问题'], row['答案']))
			self.connection.commit()
			self.logger.info("数据插入成功")
		except Exception as e:
			self.logger.error(f"数据插入失败：{e}")
			self.connection.rollback()
			raise

	def fetch_question(self):
		# 获取所有问题
		try:
			# 执行查询
			self.cursor.execute('SELECT question FROM JPKB')
			# 获取查询结果
			result = self.cursor.fetchall()
			# 记录获取成功
			self.logger.info(f"成功获取问题")
			# 返回结果
			return result
		except Exception as e:
			# 记录查询失败
			self.logger.error(f"查询失败：{e}")
			# 返回空列表
			return []

	def fetch_answer(self, question):
		# 获取指定问题的答案
		try:
			# 执行查询
			self.cursor.execute("SELECT answer FROM JPKB WHERE question=%s", (question,))
			# 获取查询结果
			result = self.cursor.fetchall()
			# 返回生成的答案或 None
			return result[0] if result else None
		except Exception as e:
			# 记录查询失败
			self.logger.error(f"查询失败：{e}")
			# 返回 None
			return None

	def close(self):
		try:
			# 关闭数据库连接
			self.connection.close()
			# 记录关闭成功
			self.logger.info("MySQL 连接已关闭")
		except Exception as e:
			# 记录关闭失败
			self.logger.error(f"MySQL 断开连接失败：{e}")


if __name__ == '__main__':
	mysql_client = MySQLClient()
	mysql_client.create_table()
	mysql_client.insert_data('../data/JP学科知识问答.csv')
